use std::cell::RefCell;

use crate::bitset::{BitSet, Idx};

pub fn primes<T: Idx>() -> Primes<T> {
    Primes::default()
}

pub fn composites<T: Idx>() -> Composites<T> {
    Composites::default()
}

thread_local! {
    static PRIMES: RefCell<Vec<usize>> = RefCell::new(Vec::new());
    static COMPOSITES: RefCell<BitSet<usize>> = RefCell::new(BitSet::new());
}

/// A struct that generates prime numbers.
pub struct Primes<T: Idx> {
    prime_index: usize,
    limit: Option<T>,
}

impl<T: Idx> Default for Primes<T> {
    fn default() -> Self {
        Primes {
            prime_index: 0,
            limit: None,
        }
    }
}

impl<T: Idx> Primes<T> {
    pub fn with_limit(mut self, limit: T) -> Self {
        self.set_limit(limit);
        self
    }
    pub fn set_limit(&mut self, limit: T) {
        self.limit = Some(limit);
    }
    pub fn is_not_composite(&self, n: T) -> bool {
        !COMPOSITES.with_borrow(|composites| composites.contains(n.index()))
    }
    fn push_prime(&mut self, prime: T) {
        PRIMES.with_borrow_mut(|primes| {
            if !primes.last().is_some_and(|&last| prime.index() <= last) {
                primes.push(prime.index());
                self.prime_index = primes.len();
            }
        });
    }
    fn set_composite(&self, composite: T) {
        COMPOSITES.with_borrow_mut(|composites| {
            composites.insert(composite.index());
        });
    }
}

impl<T: Idx> Iterator for Primes<T> {
    type Item = T;

    /// Returns the next prime number.
    fn next(&mut self) -> Option<Self::Item> {
        let next_prime = PRIMES.with_borrow(|primes| {
            if let Some(&prime) = primes.get(self.prime_index) {
                self.prime_index += 1;
                return Ok(T::new(prime));
            }
            Err(primes.last().map_or(2, |&last| last + 1))
        });
        let next_prime = match next_prime {
            Ok(prime) => return Some(prime),
            Err(next_prime) => next_prime,
        };
        for i in (next_prime..).map(T::new) {
            if self.limit.is_some_and(|limit| i > limit) {
                break;
            }
            let is_prime = self.is_not_composite(i);
            if is_prime {
                self.push_prime(i);
            }
            PRIMES.with_borrow(|primes| {
                for &p in primes {
                    let num = T::new(i.index() * p);
                    self.set_composite(num);
                    if self.limit.is_some_and(|limit| num > limit) {
                        break;
                    }
                    if i.index() % p == 0 {
                        break;
                    }
                }
            });
            if is_prime {
                return Some(i);
            }
        }
        None
    }
}

/// A struct that generates composite numbers.
pub struct Composites<T: Idx> {
    primes: Primes<T>,
    next: T,
}

impl<T: Idx> Default for Composites<T> {
    fn default() -> Self {
        Composites {
            primes: Primes::default(),
            next: T::new(4),
        }
    }
}

impl<T: Idx> Composites<T> {
    pub fn with_limit(mut self, limit: T) -> Self {
        self.set_limit(limit);
        self
    }
    pub fn set_limit(&mut self, limit: T) {
        self.primes.set_limit(limit);
    }
    pub fn is_prime(&self, n: T) -> bool {
        self.primes.is_not_composite(n)
    }
}

impl<T: Idx> Iterator for Composites<T> {
    type Item = T;

    /// Returns the next composite number.
    fn next(&mut self) -> Option<Self::Item> {
        let next = self.next;
        if self.primes.limit.is_some_and(|limit| next > limit) {
            return None;
        }
        self.next.increment_by(1);
        while let Some(prime) = self.primes.next() {
            if prime > self.next {
                break;
            }
        }
        if self.primes.is_not_composite(self.next) {
            self.next.increment_by(1);
        }
        Some(next)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_primes() {
        let primes: Vec<usize> = primes().take(10).collect();
        assert_eq!(primes, vec![2, 3, 5, 7, 11, 13, 17, 19, 23, 29]);
    }

    #[test]
    fn test_composites() {
        let composites: Vec<usize> = composites().take(10).collect();
        assert_eq!(composites, vec![4, 6, 8, 9, 10, 12, 14, 15, 16, 18]);
    }

    #[test]
    fn test_with_limits() {
        let primes: Vec<usize> = primes().with_limit(10).collect();
        assert_eq!(primes, vec![2, 3, 5, 7]);
        let composites: Vec<usize> = composites().with_limit(10).collect();
        assert_eq!(composites, vec![4, 6, 8, 9, 10]);
    }
}
